{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os,sys,inspect\n",
    "currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))\n",
    "parentdir = os.path.dirname(currentdir)\n",
    "sys.path.insert(0,parentdir) \n",
    "\n",
    "from torchnlp.datasets import penn_treebank_dataset\n",
    "import torch\n",
    "from torchnlp.samplers import BPTTBatchSampler\n",
    "from torch.utils.data import DataLoader\n",
    "from rsm_samplers import MNISTSequenceSampler, ptb_pred_sequence_collate\n",
    "from torchvision import datasets, transforms\n",
    "import torch.nn.functional as F\n",
    "from torch.nn import CrossEntropyLoss, MSELoss\n",
    "from importlib import reload \n",
    "from torch.utils.data import Sampler, BatchSampler\n",
    "import rsm\n",
    "from matplotlib.lines import Line2D\n",
    "import numpy as np\n",
    "import torchvision.utils as vutils\n",
    "from functools import reduce, partial\n",
    "import matplotlib.pyplot as plt\n",
    "from tensorboardX import SummaryWriter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import rsm_samplers\n",
    "import rsm\n",
    "import util\n",
    "import baseline_models\n",
    "reload(rsm)\n",
    "reload(rsm_samplers)\n",
    "reload(util)\n",
    "reload(baseline_models)\n",
    "\n",
    "from torch.utils.data import DataLoader, BatchSampler\n",
    "\n",
    "writer = SummaryWriter()\n",
    "\n",
    "BSZ = 10\n",
    "MB = 10\n",
    "PAGI9 = [[2, 4, 0, 7, 8, 1, 6, 1, 8], [2, 7, 4, 9, 5, 9, 3, 1, 0], [5, 7, 3, 4, 1, 3, 1, 6, 4], [1, 3, 7, 5, 2, 5, 5, 3, 4], [\n",
    "    2, 9, 1, 9, 2, 8, 3, 2, 7], [1, 2, 6, 4, 8, 3, 5, 0, 3], [3, 8, 0, 5, 6, 4, 1, 3, 9], [4, 7, 5, 3, 7, 6, 7, 2, 4]]\n",
    "\n",
    "dataset = rsm_samplers.MNISTBufferedDataset(\"~/nta/datasets\", download=True,\n",
    "                                            transform=transforms.Compose([\n",
    "                                                transforms.ToTensor(),\n",
    "                                                transforms.Normalize((0.1307,), (0.3081,))\n",
    "                                            ]),)\n",
    "\n",
    "sampler = rsm_samplers.MNISTSequenceSampler(dataset, sequences=PAGI9, \n",
    "                                            batch_size=BSZ,\n",
    "                                            max_batches=MB,\n",
    "                                            random_mnist_images=True)\n",
    "\n",
    "loader = DataLoader(dataset,\n",
    "             batch_sampler=sampler,\n",
    "             collate_fn=rsm_samplers.pred_sequence_collate)\n",
    "\n",
    "model = baseline_models.LSTMModel(\n",
    "                vocab_size=10000,\n",
    "                embed_dim=100,\n",
    "                nhid=200,\n",
    "                d_in=28**2,\n",
    "                d_out=28**2\n",
    "            )\n",
    "predictor = rsm.RSMPredictor(\n",
    "                d_in=28**2,\n",
    "                d_out=10000,\n",
    "                hidden_size=1200\n",
    "            )\n",
    "\n",
    "class BPTTTrainer():\n",
    "    def __init__(self, model, loader, k1=1, k2=30, predictor=None, bsz=BSZ):\n",
    "        self.k1 = k1\n",
    "        self.k2 = k2\n",
    "        self.model = model\n",
    "        self.loader = loader\n",
    "        self.predictor = predictor\n",
    "        self.predictor_loss = CrossEntropyLoss()\n",
    "        self.optimizer = torch.optim.Adam(params=self.model.parameters(), lr=1e-5)\n",
    "        self.pred_optimizer = torch.optim.Adam(params=self.predictor.parameters(), lr=1e-5)\n",
    "        self.bsz = bsz\n",
    "        self.loss_module = torch.nn.MSELoss()\n",
    "        self.retain_graph = self.k1 < self.k2\n",
    "        self.epoch = 0\n",
    "        \n",
    "        # Predictor counts\n",
    "        self.total_samples = 0\n",
    "        self.correct_samples = 0\n",
    "        \n",
    "        self.total_loss = 0.0\n",
    "        \n",
    "        if torch.cuda.is_available():\n",
    "            print(\"setup: Using cuda\")\n",
    "            self.device = torch.device(\"cuda\")\n",
    "            torch.cuda.manual_seed(seed)\n",
    "        else:\n",
    "            print(\"setup: Using cpu\")\n",
    "            self.device = torch.device(\"cpu\")\n",
    "            \n",
    "        self.model.to(self.device)\n",
    "        self.predictor.to(self.device)\n",
    "\n",
    "    def one_step_module(self, inp, hidden):\n",
    "        out, new_hidden = self.model(inp, hidden)\n",
    "        return (out, new_hidden)\n",
    "    \n",
    "    def _repackage_hidden(self, h):\n",
    "        \"\"\"Wraps hidden states in new Tensors, to detach them from their history.\"\"\"\n",
    "        if isinstance(h, torch.Tensor):\n",
    "            return h.detach()\n",
    "        else:\n",
    "            return tuple(self._repackage_hidden(v) for v in h)\n",
    "        \n",
    "    def predict(self, out, pred_tgts):\n",
    "        pred_out = self.predictor(out.detach())\n",
    "        loss = self.predictor_loss(pred_out, pred_tgts)\n",
    "        loss.backward()\n",
    "        self.pred_optimizer.step()\n",
    "        _, class_predictions = torch.max(pred_out, 1)\n",
    "        self.total_samples += pred_tgts.size(0)\n",
    "        correct_arr = class_predictions == pred_tgts\n",
    "        self.correct_samples += correct_arr.sum().item()        \n",
    "        \n",
    "    def train(self):\n",
    "        self.total_samples = 0\n",
    "        self.correct_samples = 0\n",
    "        self.total_loss = 0.0\n",
    "\n",
    "        states = [(None, self.model.init_hidden(self.bsz))]\n",
    "\n",
    "        outputs = []\n",
    "        targets = []\n",
    "\n",
    "        for i, (inp, target, pred_tgts, input_labels) in enumerate(self.loader):\n",
    "            inp = inp.to(self.device)\n",
    "            target = target.to(self.device)\n",
    "            pred_tgts = pred_tgts.to(self.device)\n",
    "\n",
    "            batch_loss = 0.0\n",
    "            state = self._repackage_hidden(states[-1][1])\n",
    "            for h in state:\n",
    "                h.requires_grad=True\n",
    "            output, new_state = self.one_step_module(inp, state)\n",
    "\n",
    "            outputs.append(output)\n",
    "            targets.append(target)\n",
    "            while len(outputs) > self.k1:\n",
    "                # Delete stuff that is too old\n",
    "                del outputs[0]\n",
    "                del targets[0]\n",
    "\n",
    "            states.append((state, new_state))\n",
    "            while len(states) > self.k2:\n",
    "                # Delete stuff that is too old\n",
    "                del states[0]\n",
    "                \n",
    "            if (i+1)%self.k1 == 0:\n",
    "                self.optimizer.zero_grad()\n",
    "                # backprop last module (keep graph only if they ever overlap)\n",
    "                start = time.time()\n",
    "                for j in range(self.k2-1):\n",
    "                    # print('j', j)\n",
    "                    if j < self.k1:\n",
    "                        loss = self.loss_module(outputs[-j-1], targets[-j-1])\n",
    "                        batch_loss += loss.item()\n",
    "                        loss.backward(retain_graph=True)\n",
    "\n",
    "                    # if we get all the way back to the \"init_state\", stop\n",
    "                    if states[-j-2][0] is None:\n",
    "                        break\n",
    "                    curr_h_grad = states[-j-1][0][0].grad\n",
    "                    curr_c_grad = states[-j-1][0][1].grad                    \n",
    "                    states[-j-2][1][0].backward(curr_h_grad, retain_graph=self.retain_graph)\n",
    "                    states[-j-2][1][1].backward(curr_c_grad, retain_graph=self.retain_graph)                    \n",
    "                # print(\"opt step, batch loss: %.3f\" % batch_loss)\n",
    "                self.optimizer.step()\n",
    "                self.total_loss += batch_loss\n",
    "            \n",
    "            self.predict(output, pred_tgts)\n",
    "        \n",
    "        if self.total_samples:\n",
    "            train_acc = 100.*self.correct_samples/self.total_samples\n",
    "            print(self.epoch, \"train acc: %.3f%%, loss: %.3f\" % (train_acc, self.total_loss))\n",
    "            writer.add_scalar('train_acc', train_acc, self.epoch)\n",
    "            writer.add_scalar('train_loss', self.total_loss, self.epoch)            \n",
    "        \n",
    "        self.epoch += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pred acc: 0.400%, loss: 50.800\n"
     ]
    }
   ],
   "source": [
    "trainer = BPTTTrainer(model, loader, predictor=predictor, k1=1, k2=30)\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([5, 5])\n",
      "tensor([6, 2])\n",
      "tensor([4, 5])\n",
      "tensor([1, 5])\n",
      "tensor([3, 3])\n",
      "tensor([9, 4])\n",
      "tensor([2, 3])\n",
      "tensor([9, 8])\n",
      "tensor([1, 0])\n",
      "tensor([9, 5])\n",
      "new epoch\n",
      "tensor([2, 6])\n",
      "tensor([8, 4])\n",
      "tensor([3, 1])\n",
      "tensor([2, 3])\n",
      "tensor([7, 9])\n",
      "tensor([2, 1])\n",
      "tensor([4, 2])\n",
      "tensor([0, 6])\n",
      "tensor([7, 4])\n",
      "tensor([8, 8])\n"
     ]
    }
   ],
   "source": [
    "# PAGI9 = [[2, 4, 0, 7, 8, 1, 6, 1, 8], [2, 7, 4, 9, 5, 9, 3, 1, 0], [5, 7, 3, 4, 1, 3, 1, 6, 4], [1, 3, 7, 5, 2, 5, 5, 3, 4], [\n",
    "#    2, 9, 1, 9, 2, 8, 3, 2, 7], [1, 2, 6, 4, 8, 3, 5, 0, 3], [3, 8, 0, 5, 6, 4, 1, 3, 9], [4, 7, 5, 3, 7, 6, 7, 2, 4]]\n",
    "\n",
    "for i, (inp, target, pred_tgts, input_labels) in enumerate(loader):\n",
    "    print(input_labels[0:2])\n",
    "\n",
    "print('new epoch')\n",
    "for i, (inp, target, pred_tgts, input_labels) in enumerate(loader):\n",
    "    print(input_labels[0:2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Samples\n",
    "\n",
    "#### From https://discuss.pytorch.org/t/truncated-backprop-data-clarification/34854"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.5388, 0.3831, 0.0448, 0.9141, 0.7032],\n",
      "        [0.4118, 0.6417, 0.1732, 0.8667, 0.4823],\n",
      "        [0.5891, 0.5699, 0.9462, 0.2595, 0.5103]])\n",
      "tensor([[2.5839],\n",
      "        [2.5757],\n",
      "        [2.8750]])\n"
     ]
    }
   ],
   "source": [
    "a = torch.rand(3, 5)\n",
    "print(a)\n",
    "\n",
    "sums = a.sum(dim=1, keepdim=True)\n",
    "print(sums)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.6000, 1.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.0000, 0.6000, 1.0000, 0.0000, 0.0000],\n",
      "        [0.0000, 1.0000, 0.0000, 0.6000, 0.0000]])\n",
      "tensor([[0.3750, 0.6250, 0.0000, 0.0000, 0.0000],\n",
      "        [0.0000, 0.3750, 0.6250, 0.0000, 0.0000],\n",
      "        [0.0000, 0.6250, 0.0000, 0.3750, 0.0000]])\n"
     ]
    }
   ],
   "source": [
    "cache = torch.zeros(3, 5)\n",
    "index = torch.tensor([0, 1, 3])\n",
    "cache.scatter_(1, index.unsqueeze(1), 1.0)\n",
    "\n",
    "cache = cache * 0.6\n",
    "\n",
    "cache.scatter_(1, torch.tensor([1, 2, 1]).unsqueeze(1), 1.0)\n",
    "print(cache)\n",
    "\n",
    "print(cache / cache.sum(dim=1, keepdim=True))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
